# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Core engine of MindSPONGE
"""

import torch
import os
from typing import Union, List
import time
import datetime
from collections.abc import Iterable

from torch import nn, Tensor
from torch.optim import Optimizer

from .simulation import WithEnergyCell, WithForceCell
from .simulation import RunOneStepCell
from ..function import any_not_none, get_arguments
from ..potential import PotentialCell, ForceCell
from ..optimizer import Updater, UpdaterMD
from ..system import Molecule
from ..callback import print_info, WriteH5MD
# from ..metrics import MetricCV, get_metrics


class Sponge:

    def __init__(self,
                 network: Union[Molecule, WithEnergyCell, RunOneStepCell],
                 potential: PotentialCell = None,
                 optimizer: Optimizer = None,
                 metrics: dict = None,
                 write_h5md: WriteH5MD = None,
                 **kwargs
                 ):
        self._kwargs = get_arguments(locals(), kwargs)

        self._create_time = int(time.time() * 1e9)

        self._force_function = None
        if optimizer is None:
            if potential is not None:
                raise ValueError('When optimizer is None, potential must also be None!')
            # network is RunOneStepCell: Sponge = RunOneStepCell
            self._simulation_network: RunOneStepCell = network
            self._system_with_energy: WithEnergyCell = self._simulation_network.system_with_energy
            self._system_with_force: WithForceCell = self._simulation_network.system_with_force
            self._optimizer: Optimizer = self._simulation_network.optimizer
            self._system: Molecule = self._simulation_network.system

            self._potential_function = None
            if self._system_with_energy is not None:
                self._potential_function: PotentialCell = self._system_with_energy.potential_function

            if self._system_with_force is not None:
                self._force_function: ForceCell = self._system_with_force.force_function
        else:
            self._system_with_force = None
            self._optimizer = optimizer
            if potential is None:
                # network is WithEnergyCell: Sponge = WithEnergyCell + optimizer
                self._system_with_energy: WithEnergyCell = network
                self._simulation_network = RunOneStepCell(
                    energy=self._system_with_energy, optimizer=self._optimizer)
                self._system: Molecule = self._system_with_energy.system
                self._potential_function: PotentialCell = self._system_with_energy.potential_function
            else:
                # network is system: Sponge = system + potential + optimizer
                self._system: Molecule = network
                self._potential_function: PotentialCell = potential
                self._system_with_energy = WithEnergyCell(self._system, self._potential_function)
                self._simulation_network = RunOneStepCell(
                    energy=self._system_with_energy, optimizer=self._optimizer)

        self._metrics = metrics

        self._num_energies = self._simulation_network.num_energies
        self._num_biases = self._simulation_network.num_biases
        self._energy_names = self._simulation_network.energy_names

        self.use_updater = False
        if isinstance(self._optimizer, Updater):
            self.use_updater = True

        if isinstance(self._simulation_network, RunOneStepCell):
            self._simulation_network.set_pbc_grad(self.use_updater)

        self.units = self._system.units
        if isinstance(self._optimizer, UpdaterMD):
            self.time_step = self._optimizer.time_step
        else:
            self.time_step = self._optimizer.param_groups[0]['lr']

        self.coordinate = self._system.coordinate
        self.pbc_box = self._system.pbc_box

        self.energy_neighbour_list = None
        if self._system_with_energy is not None:
            self.energy_neighbour_list = self._system_with_energy.neighbour_list

        self.force_neighbour_list = None
        if self._system_with_force is not None:
            self.force_neighbour_list = self._system_with_force.neighbour_list

        self.neighbour_list_pace = self._simulation_network.neighbour_list_pace

        self.energy_cutoff = self._simulation_network.energy_cutoff
        self.force_cutoff = self._simulation_network.force_cutoff

        self.update_nl = any_not_none([self.energy_cutoff, self.force_cutoff])

        self._metric_value = None
        if metrics is not None:
            self._metric_value = {}
            for k, v in self._metrics.items():
                self._metric_value[k] = 0


        self.sim_step = 0
        self.sim_time = 0.0

        self._potential = None
        self._force = None
        self._bias = None

        self._use_bias = False
        self.write_h5md = write_h5md

    def print_info(self):
        print(f"{'--'*40}")
        if self._system_with_energy is not None: self._system_with_energy.print_info()
        if self._system_with_force is not None: self._system_with_force.print_info()
        if self._optimizer is not None:
            print("[TorchSPONGE] Parameter set for Updater is:")
            if isinstance(self._optimizer, UpdaterMD):
                if self._optimizer.integrator.name == 'PILE':
                    print(f"[TorchSPONGE] Type: PIMD")
                    print(f"[TorchSPONGE] Bead_num: {self._optimizer.integrator.num_walker}")
                    print(f"[TorchSPONGE] Ensemble: NVT")
                    print(f"[TorchSPONGE] Temperature: {self._optimizer.integrator.temperature} K")
                    print(f"[TorchSPONGE] Thermostat: Langevin [gamma:{self._optimizer.integrator.gamma0}]")
                else:
                    print(f"[TorchSPONGE] Type: MD")
                    if self._optimizer.thermostat is not None and self._optimizer.barostat is not None:
                        print(f"[TorchSPONGE] Ensemble: NPT")
                        print(f"[TorchSPONGE] Temperature: {self._optimizer.thermostat.temperature} K")
                        print(f"[TorchSPONGE] Pressure: {self._optimizer.barostat.pressure} Bar")
                    if self._optimizer.thermostat is None and self._optimizer.barostat is None:
                        print(f"[TorchSPONGE] Ensemble: NVE")
                    if self._optimizer.thermostat is not None and self._optimizer.barostat is None:
                        print(f"[TorchSPONGE] Ensemble: NVT")
                        print(f"[TorchSPONGE] Temperature: {self._optimizer.thermostat.temperature} K")
                    if self._optimizer.thermostat is None and self._optimizer.barostat is not None:
                        print(f"[TorchSPONGE] Ensemble: NPH")
                        print(f"[TorchSPONGE] Pressure: {self._optimizer.barostat.pressure} Bar")
                print(f"[TorchSPONGE] Time_step: {self.time_step:.6f} ps")
            elif isinstance(self._optimizer, torch.optim.Optimizer):
                print(f"[TorchSPONGE] Type: Minimization")
                print(f"[TorchSPONGE] Learning_rate: {self.time_step}")
            print(f"{'--'*40}")
        
        if self.write_h5md is not None:
            print(f"[TorchSPONGE] Trajectory is saved to {self.write_h5md.filename}")
            print(f"{'--'*40}")
    
    @property
    def energy_names(self) -> List[str]:
        """names of energy terms

        Return:
            list of str, names of energy terms

        """
        return self._energy_names

    @property
    def num_energies(self) -> int:
        """number of energy terms

        Return:
            int, number of energy terms

        """
        return self._num_energies

    @property
    def num_biases(self) -> int:
        """number of bias potential energies V

        Return:
            int, number of bias potential energies

        """
        return self._num_biases
    
    @property
    def bias_names(self) -> List[str]:
        """names of energy terms

        Return:
            list of str, names of energy terms

        """
        if self.num_biases > 0:
            return self._simulation_network.bias_names
        
        return None

    def update_neighbour_list(self):
        """update neighbour list"""
        self._simulation_network.update_neighbour_list()
        return self

    def update_bias(self, step: int):
        """update bias potential"""
        self._simulation_network.update_bias(step)

    def update_wrapper(self, step: int):
        """update energy wrapper"""
        self._simulation_network.update_wrapper(step)

    def update_modifier(self, step: int):
        """update force modifier"""
        self._simulation_network.update_modifier(step)

    def change_optimizer(self, optimizer: Optimizer):
        """
        change optimizer.
        """
        if self._optimizer is None:
            raise ValueError('Cannot change the optimizer, because the initial optimizer is None '
                             'or the network is not a RunOneStepCell type.')

        self._optimizer = optimizer

        if isinstance(self._optimizer, Updater):
            self.use_updater = True
        else:
            self.use_updater = False

        self._simulation_network = RunOneStepCell(
            energy=self._system_with_energy, optimizer=self._optimizer)
        self._simulation_network.set_pbc_grad(self.use_updater)

        return self

    def change_potential(self, potential: PotentialCell):
        """
        change potential energy.
        """
        if self._potential_function is None:
            raise ValueError('Cannot change the potential, because the initial potential is None '
                             'or the network is not a WithEnergyCell type.')
        if self._optimizer is None:
            raise ValueError('Cannot change the potential, because the initial optimizer is None '
                             'or the network is not a RunOneStepCell type.')

        self._potential_function = potential
        self._system_with_energy = WithEnergyCell(self._system, self._potential_function)
        self._simulation_network = RunOneStepCell(
            energy=self._system_with_energy, optimizer=self._optimizer)
        self._simulation_network.set_pbc_grad(self.use_updater)

        return self

    def calc_energy(self) -> Tensor:
        """calculate the total potential energy (potential energy and bias potential) of the simulation system.

        Return:
            energy (Tensor):    Tensor of shape `(B, 1)`. Data type is float.
                                Total potential energy.

        """
        if self._system_with_energy is None:
            return None
        return self._system_with_energy()

    def calc_energies(self) -> Tensor:
        """calculate the energy terms of the potential energy.

        Return:
            energies (Tensor):  Tensor of shape `(B, U)`. Data type is float.
                                Energy terms.

        Symbols:
            B:  Batchsize, i.e. number of walkers of the simulation.
            U:  Number of potential energy terms.

        """
        if self._system_with_energy is None:
            return None
        return self._system_with_energy.calc_energies()

    def calc_biases(self) -> Tensor:
        """calculate the bias potential terms.

        Return:
            biases (Tensor):    Tensor of shape `(B, V)`. Data type is float.
                                Energy terms.

        Symbols:
            B:  Batchsize, i.e. number of walkers of the simulation.
            V:  Number of bias potential terms.

        """
        if self._system_with_energy is None:
            return None
        return self._system_with_energy.calc_biases()
    

    def run(self,
            steps: int,
            print_interval: int = 100,
            write_interval: int = 100,
            ):
        """Simulation API.
        """

        beg_time = datetime.datetime.now()
        self.print_info()
        print('[TorchSPONGE] Started simulation at', beg_time.strftime('%Y-%m-%d %H:%M:%S'))
        self._simulation_process(steps,print_interval,write_interval)
        end_time = datetime.datetime.now()

        print('[TorchSPONGE] Finished simulation at', end_time.strftime('%Y-%m-%d %H:%M:%S'))
        used_time = end_time - beg_time
        d = used_time.days
        s = used_time.seconds
        m, s = divmod(s, 60)
        h, m = divmod(m, 60)
        if d >= 1:
            print('[TorchSPONGE] Simulation time: %d days, %d hours, %d minutes and %d seconds.' % (d, h, m, s))
        elif h >= 1:
            print('[TorchSPONGE] Simulation time: %d hours %d minutes %d seconds.' % (h, m, s))
        elif m >= 1:
            s += used_time.microseconds / 1e6
            print('[TorchSPONGE] Simulation time: %d minutes %1.1f seconds.' % (m, s))
        else:
            s += used_time.microseconds / 1e6
            print('[TorchSPONGE] Simulation time: %1.2f seconds.' % s)
        print('-'*80)
        return self

    def calc_potential(self) -> Tensor:
        """calculate and return the potential energy"""
        if self._system_with_energy is None:
            return 0
        return self._system_with_energy()

    def get_energies(self) -> Tensor:
        """Tensor of potential energies with shape (B, U). Data type is float."""
        return self._simulation_network.energies

    def get_biases(self) -> Tensor:
        """Tensor bias potential energies with shape (B, V). Data type is float."""
        return self._simulation_network.biases

    def get_bias(self) -> Tensor:
        """Tensor of total bias potential with shape (B, 1). Data type is float."""
        return self._simulation_network.bias
    def get_meta_potential(self):
        if self._simulation_network.bias_function is None:
            return None,None
        else:
            hills = [i.hills.detach().cpu().numpy() for i in self._simulation_network.bias_function]
            grids = [i.grids.detach().cpu().numpy() for i in self._simulation_network.bias_function]
            return grids, hills

    def _simulation_process(self,
                            steps: int,
                            print_interval: int,
                            write_interval: int,
                            ):
        """
        Training process. The data would be passed to network directly.
        """

        self.sim_step = steps
        self.sim_time = steps 
        grids, hills = self.get_meta_potential()
        if self.write_h5md is not None: self.write_h5md.begin(self.num_energies, self.energy_names, self.num_biases, self.bias_names, self._metrics, grids, hills)
     
        for i in range(steps):
            if self.neighbour_list_pace > 0 and i % self.neighbour_list_pace == 0:
                self.update_neighbour_list()
            self._run_one_step(i)
            if i % print_interval == 0:
                if self._metrics is not None:
                    for k, v in self._metrics.items():
                        self._metric_value[k] = v(self._system.coordinate, self._system.pbc_box)
                print_info(i, self._potential, self._bias, self._optimizer, self._metric_value)
            if write_interval <= 0:
                continue
            elif (i % write_interval == 0 or i == steps-1) and self.write_h5md is not None:
                if self._metrics is not None:
                    for k, v in self._metrics.items():
                        self._metric_value[k] = v(self._system.coordinate, self._system.pbc_box)
                pressure = self._optimizer.pressure.sum(-1) if self._optimizer.pressure is not None else None
                grids, hills = self.get_meta_potential()
                self.write_h5md.write_info(i, i * self.time_step, self.coordinate, self._optimizer.velocity, self._force, self._potential,
                                           self.get_energies(), self._optimizer.kinetics.sum(-1), self._optimizer.temperature, self.pbc_box, 
                                           pressure, self.get_bias(), self.get_biases(), hills, self._metric_value)
                
        if self.write_h5md is not None: 
            grids, hills = self.get_meta_potential()
            self.write_h5md.end(hills)

    def _run_one_step(self,
                      steps: int,
                       ):
        """run one epoch simulation"""

        self._potential, self._force = self._simulation_network()
        self._bias = self.calc_biases()
        self.update_bias(steps)
        self.update_wrapper(steps)
        self.update_modifier(steps)

        return 0

    # def _clear_metrics(self):
    #     """Clear metrics local values."""
    #     for metric in self._metric_fns.values():
    #         metric.clear()

    # def _update_metrics(self, outputs):
    #     """Update metrics local values."""
    #     if isinstance(outputs, Tensor):
    #         outputs = (outputs,)
    #     if not isinstance(outputs, tuple):
    #         raise ValueError(
    #             f"The argument 'outputs' should be tuple, but got {type(outputs)}.")

    #     for metric in self._metric_fns.values():
    #         metric.update(*outputs)

    # def _get_metrics(self):
    #     """Get metrics local values."""
    #     metrics = dict()
    #     for key, value in self._metric_fns.items():
    #         metrics[key] = value.eval()
    #     return metrics

    @property
    def create_time(self):
        return self._create_time